{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Convolutional GANs\n",
    "\n",
    "In this notebook, you'll build a GAN using convolutional layers in the generator and discriminator. This is called a Deep Convolutional GAN, or DCGAN for short. The DCGAN architecture was first explored last year and has seen impressive results in generating new images, you can read the [original paper here](https://arxiv.org/pdf/1511.06434.pdf).\n",
    "\n",
    "You'll be training DCGAN on the [Street View House Numbers](http://ufldl.stanford.edu/housenumbers/) (SVHN) dataset. These are color images of house numbers collected from Google street view. SVHN images are in color and much more variable than MNIST. \n",
    "\n",
    "![SVHN Examples](assets/SVHN_examples.png)\n",
    "\n",
    "So, we'll need a deeper and more powerful network. This is accomplished through using convolutional layers in the discriminator and generator. It's also necessary to use batch normalization to get the convolutional networks to train. The only real changes compared to what [you saw previously](https://github.com/udacity/deep-learning/tree/master/gan_mnist) are in the generator and discriminator, otherwise the rest of the implementation is the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import pickle as pkl\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.io import loadmat\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "!mkdir data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the data\n",
    "\n",
    "Here you can download the SVHN dataset. Run the cell above and it'll download to your machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from urllib.request import urlretrieve\n",
    "from os.path import isfile, isdir\n",
    "from tqdm import tqdm\n",
    "\n",
    "data_dir = 'data/'\n",
    "\n",
    "if not isdir(data_dir):\n",
    "    raise Exception(\"Data directory doesn't exist!\")\n",
    "\n",
    "class DLProgress(tqdm):\n",
    "    last_block = 0\n",
    "\n",
    "    def hook(self, block_num=1, block_size=1, total_size=None):\n",
    "        self.total = total_size\n",
    "        self.update((block_num - self.last_block) * block_size)\n",
    "        self.last_block = block_num\n",
    "\n",
    "if not isfile(data_dir + \"train_32x32.mat\"):\n",
    "    with DLProgress(unit='B', unit_scale=True, miniters=1, desc='SVHN Training Set') as pbar:\n",
    "        urlretrieve(\n",
    "            'http://ufldl.stanford.edu/housenumbers/train_32x32.mat',\n",
    "            data_dir + 'train_32x32.mat',\n",
    "            pbar.hook)\n",
    "\n",
    "if not isfile(data_dir + \"test_32x32.mat\"):\n",
    "    with DLProgress(unit='B', unit_scale=True, miniters=1, desc='SVHN Training Set') as pbar:\n",
    "        urlretrieve(\n",
    "            'http://ufldl.stanford.edu/housenumbers/test_32x32.mat',\n",
    "            data_dir + 'test_32x32.mat',\n",
    "            pbar.hook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These SVHN files are `.mat` files typically used with Matlab. However, we can load them in with `scipy.io.loadmat` which we imported above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "trainset = loadmat(data_dir + 'train_32x32.mat')\n",
    "testset = loadmat(data_dir + 'test_32x32.mat')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I'm showing a small sample of the images. Each of these is 32x32 with 3 color channels (RGB). These are the real images we'll pass to the discriminator and what the generator will eventually fake."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "idx = np.random.randint(0, trainset['X'].shape[3], size=36)\n",
    "fig, axes = plt.subplots(6, 6, sharex=True, sharey=True, figsize=(5,5),)\n",
    "for ii, ax in zip(idx, axes.flatten()):\n",
    "    ax.imshow(trainset['X'][:,:,:,ii], aspect='equal')\n",
    "    ax.xaxis.set_visible(False)\n",
    "    ax.yaxis.set_visible(False)\n",
    "plt.subplots_adjust(wspace=0, hspace=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we need to do a bit of preprocessing and getting the images into a form where we can pass batches to the network. First off, we need to rescale the images to a range of -1 to 1, since the output of our generator is also in that range. We also have a set of test and validation images which could be used if we're trying to identify the numbers in the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def scale(x, feature_range=(-1, 1)):\n",
    "    # scale to (0, 1)\n",
    "    x = ((x - x.min())/(255 - x.min()))\n",
    "    \n",
    "    # scale to feature_range\n",
    "    min, max = feature_range\n",
    "    x = x * (max - min) + min\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Dataset:\n",
    "    def __init__(self, train, test, val_frac=0.5, shuffle=False, scale_func=None):\n",
    "        split_idx = int(len(test['y'])*(1 - val_frac))\n",
    "        self.test_x, self.valid_x = test['X'][:,:,:,:split_idx], test['X'][:,:,:,split_idx:]\n",
    "        self.test_y, self.valid_y = test['y'][:split_idx], test['y'][split_idx:]\n",
    "        self.train_x, self.train_y = train['X'], train['y']\n",
    "        \n",
    "        self.train_x = np.rollaxis(self.train_x, 3)\n",
    "        self.valid_x = np.rollaxis(self.valid_x, 3)\n",
    "        self.test_x = np.rollaxis(self.test_x, 3)\n",
    "        \n",
    "        if scale_func is None:\n",
    "            self.scaler = scale\n",
    "        else:\n",
    "            self.scaler = scale_func\n",
    "        self.shuffle = shuffle\n",
    "        \n",
    "    def batches(self, batch_size):\n",
    "        if self.shuffle:\n",
    "            idx = np.arange(len(dataset.train_x))\n",
    "            np.random.shuffle(idx)\n",
    "            self.train_x = self.train_x[idx]\n",
    "            self.train_y = self.train_y[idx]\n",
    "        \n",
    "        n_batches = len(self.train_y)//batch_size\n",
    "        for ii in range(0, len(self.train_y), batch_size):\n",
    "            x = self.train_x[ii:ii+batch_size]\n",
    "            y = self.train_y[ii:ii+batch_size]\n",
    "            \n",
    "            yield self.scaler(x), y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network Inputs\n",
    "\n",
    "Here, just creating some placeholders like normal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def model_inputs(real_dim, z_dim):\n",
    "    inputs_real = tf.placeholder(tf.float32, (None, *real_dim), name='input_real')\n",
    "    inputs_z = tf.placeholder(tf.float32, (None, z_dim), name='input_z')\n",
    "    \n",
    "    return inputs_real, inputs_z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generator\n",
    "\n",
    "Here you'll build the generator network. The input will be our noise vector `z` as before. Also as before, the output will be a $tanh$ output, but this time with size 32x32 which is the size of our SVHN images.\n",
    "\n",
    "What's new here is we'll use convolutional layers to create our new images. The first layer is a fully connected layer which is reshaped into a deep and narrow layer, something like 4x4x1024 as in the original DCGAN paper. Then we use batch normalization and a leaky ReLU activation. Next is a transposed convolution where typically you'd halve the depth and double the width and height of the previous layer. Again, we use batch normalization and leaky ReLU. For each of these layers, the general scheme is convolution > batch norm > leaky ReLU.\n",
    "\n",
    "You keep stacking layers up like this until you get the final transposed convolution layer with shape 32x32x3. Below is the archicture used in the original DCGAN paper:\n",
    "\n",
    "![DCGAN Generator](assets/dcgan.png)\n",
    "\n",
    "Note that the final layer here is 64x64x3, while for our SVHN dataset, we only want it to be 32x32x3. \n",
    "\n",
    ">**Exercise:** Build the transposed convolutional network for the generator in the function below. Be sure to use leaky ReLUs on all the layers except for the last tanh layer, as well as batch normalization on all the transposed convolutional layers except the last one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def generator(z, output_dim, reuse=False, alpha=0.2, training=True):\n",
    "    with tf.variable_scope('generator', reuse=reuse):\n",
    "        # First fully connected layer\n",
    "        x\n",
    "        \n",
    "        # Output layer, 32x32x3\n",
    "        logits = \n",
    "        \n",
    "        out = tf.tanh(logits)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminator\n",
    "\n",
    "Here you'll build the discriminator. This is basically just a convolutional classifier like you've build before. The input to the discriminator are 32x32x3 tensors/images. You'll want a few convolutional layers, then a fully connected layer for the output. As before, we want a sigmoid output, and you'll need to return the logits as well. For the depths of the convolutional layers I suggest starting with 16, 32, 64 filters in the first layer, then double the depth as you add layers. Note that in the DCGAN paper, they did all the downsampling using only strided convolutional layers with no maxpool layers.\n",
    "\n",
    "You'll also want to use batch normalization with `tf.layers.batch_normalization` on each layer except the first convolutional and output layers. Again, each layer should look something like convolution > batch norm > leaky ReLU.\n",
    "\n",
    "Note: in this project, your batch normalization layers will always use batch statistics. (That is, always set `training` to `True`.) That's because we are only interested in using the discriminator to help train the generator. However, if you wanted to use the discriminator for inference later, then you would need to set the `training` parameter appropriately.\n",
    "\n",
    ">**Exercise:** Build the convolutional network for the discriminator. The input is a 32x32x3 images, the output is a sigmoid plus the logits. Again, use Leaky ReLU activations and batch normalization on all the layers except the first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def discriminator(x, reuse=False, alpha=0.2):\n",
    "    with tf.variable_scope('discriminator', reuse=reuse):\n",
    "        # Input layer is 32x32x3\n",
    "        x =\n",
    "        \n",
    "        logits = \n",
    "        out = \n",
    "        \n",
    "        return out, logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Loss\n",
    "\n",
    "Calculating the loss like before, nothing new here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def model_loss(input_real, input_z, output_dim, alpha=0.2):\n",
    "    \"\"\"\n",
    "    Get the loss for the discriminator and generator\n",
    "    :param input_real: Images from the real dataset\n",
    "    :param input_z: Z input\n",
    "    :param out_channel_dim: The number of channels in the output image\n",
    "    :return: A tuple of (discriminator loss, generator loss)\n",
    "    \"\"\"\n",
    "    g_model = generator(input_z, output_dim, alpha=alpha)\n",
    "    d_model_real, d_logits_real = discriminator(input_real, alpha=alpha)\n",
    "    d_model_fake, d_logits_fake = discriminator(g_model, reuse=True, alpha=alpha)\n",
    "\n",
    "    d_loss_real = tf.reduce_mean(\n",
    "        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, labels=tf.ones_like(d_model_real)))\n",
    "    d_loss_fake = tf.reduce_mean(\n",
    "        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.zeros_like(d_model_fake)))\n",
    "    g_loss = tf.reduce_mean(\n",
    "        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_model_fake)))\n",
    "\n",
    "    d_loss = d_loss_real + d_loss_fake\n",
    "\n",
    "    return d_loss, g_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizers\n",
    "\n",
    "Not much new here, but notice how the train operations are wrapped in a `with tf.control_dependencies` block so the batch normalization layers can update their population statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def model_opt(d_loss, g_loss, learning_rate, beta1):\n",
    "    \"\"\"\n",
    "    Get optimization operations\n",
    "    :param d_loss: Discriminator loss Tensor\n",
    "    :param g_loss: Generator loss Tensor\n",
    "    :param learning_rate: Learning Rate Placeholder\n",
    "    :param beta1: The exponential decay rate for the 1st moment in the optimizer\n",
    "    :return: A tuple of (discriminator training operation, generator training operation)\n",
    "    \"\"\"\n",
    "    # Get weights and bias to update\n",
    "    t_vars = tf.trainable_variables()\n",
    "    d_vars = [var for var in t_vars if var.name.startswith('discriminator')]\n",
    "    g_vars = [var for var in t_vars if var.name.startswith('generator')]\n",
    "\n",
    "    # Optimize\n",
    "    with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):\n",
    "        d_train_opt = tf.train.AdamOptimizer(learning_rate, beta1=beta1).minimize(d_loss, var_list=d_vars)\n",
    "        g_train_opt = tf.train.AdamOptimizer(learning_rate, beta1=beta1).minimize(g_loss, var_list=g_vars)\n",
    "\n",
    "    return d_train_opt, g_train_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the model\n",
    "\n",
    "Here we can use the functions we defined about to build the model as a class. This will make it easier to move the network around in our code since the nodes and operations in the graph are packaged in one object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class GAN:\n",
    "    def __init__(self, real_size, z_size, learning_rate, alpha=0.2, beta1=0.5):\n",
    "        tf.reset_default_graph()\n",
    "        \n",
    "        self.input_real, self.input_z = model_inputs(real_size, z_size)\n",
    "        \n",
    "        self.d_loss, self.g_loss = model_loss(self.input_real, self.input_z,\n",
    "                                              real_size[2], alpha=0.2)\n",
    "        \n",
    "        self.d_opt, self.g_opt = model_opt(self.d_loss, self.g_loss, learning_rate, beta1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a function for displaying generated images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def view_samples(epoch, samples, nrows, ncols, figsize=(5,5)):\n",
    "    fig, axes = plt.subplots(figsize=figsize, nrows=nrows, ncols=ncols, \n",
    "                             sharey=True, sharex=True)\n",
    "    for ax, img in zip(axes.flatten(), samples[epoch]):\n",
    "        ax.axis('off')\n",
    "        img = ((img - img.min())*255 / (img.max() - img.min())).astype(np.uint8)\n",
    "        ax.set_adjustable('box-forced')\n",
    "        im = ax.imshow(img, aspect='equal')\n",
    "   \n",
    "    plt.subplots_adjust(wspace=0, hspace=0)\n",
    "    return fig, axes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And another function we can use to train our network. Notice when we call `generator` to create the samples to display, we set `training` to `False`. That's so the batch normalization layers will use the population statistics rather than the batch statistics. Also notice that we set the `net.input_real` placeholder when we run the generator's optimizer. The generator doesn't actually use it, but we'd get an errror without it because of the `tf.control_dependencies` block we created in `model_opt`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train(net, dataset, epochs, batch_size, print_every=10, show_every=100, figsize=(5,5)):\n",
    "    saver = tf.train.Saver()\n",
    "    sample_z = np.random.uniform(-1, 1, size=(72, z_size))\n",
    "\n",
    "    samples, losses = [], []\n",
    "    steps = 0\n",
    "\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        for e in range(epochs):\n",
    "            for x, y in dataset.batches(batch_size):\n",
    "                steps += 1\n",
    "\n",
    "                # Sample random noise for G\n",
    "                batch_z = np.random.uniform(-1, 1, size=(batch_size, z_size))\n",
    "\n",
    "                # Run optimizers\n",
    "                _ = sess.run(net.d_opt, feed_dict={net.input_real: x, net.input_z: batch_z})\n",
    "                _ = sess.run(net.g_opt, feed_dict={net.input_z: batch_z, net.input_real: x})\n",
    "\n",
    "                if steps % print_every == 0:\n",
    "                    # At the end of each epoch, get the losses and print them out\n",
    "                    train_loss_d = net.d_loss.eval({net.input_z: batch_z, net.input_real: x})\n",
    "                    train_loss_g = net.g_loss.eval({net.input_z: batch_z})\n",
    "\n",
    "                    print(\"Epoch {}/{}...\".format(e+1, epochs),\n",
    "                          \"Discriminator Loss: {:.4f}...\".format(train_loss_d),\n",
    "                          \"Generator Loss: {:.4f}\".format(train_loss_g))\n",
    "                    # Save losses to view after training\n",
    "                    losses.append((train_loss_d, train_loss_g))\n",
    "\n",
    "                if steps % show_every == 0:\n",
    "                    gen_samples = sess.run(\n",
    "                                   generator(net.input_z, 3, reuse=True, training=False),\n",
    "                                   feed_dict={net.input_z: sample_z})\n",
    "                    samples.append(gen_samples)\n",
    "                    _ = view_samples(-1, samples, 6, 12, figsize=figsize)\n",
    "                    plt.show()\n",
    "\n",
    "        saver.save(sess, './checkpoints/generator.ckpt')\n",
    "\n",
    "    with open('samples.pkl', 'wb') as f:\n",
    "        pkl.dump(samples, f)\n",
    "    \n",
    "    return losses, samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameters\n",
    "\n",
    "GANs are very senstive to hyperparameters. A lot of experimentation goes into finding the best hyperparameters such that the generator and discriminator don't overpower each other. Try out your own hyperparameters or read [the DCGAN paper](https://arxiv.org/pdf/1511.06434.pdf) to see what worked for them.\n",
    "\n",
    ">**Exercise:** Find hyperparameters to train this GAN. The values found in the DCGAN paper work well, or you can experiment on your own. In general, you want the discriminator loss to be around 0.3, this means it is correctly classifying images as fake or real about 50% of the time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "real_size = (32,32,3)\n",
    "z_size = 100\n",
    "learning_rate = 0.001\n",
    "batch_size = 64\n",
    "epochs = 1\n",
    "alpha = 0.01\n",
    "beta1 = 0.9\n",
    "\n",
    "# Create the network\n",
    "net = GAN(real_size, z_size, learning_rate, alpha=alpha, beta1=beta1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Load the data and train the network here\n",
    "dataset = Dataset(trainset, testset)\n",
    "losses, samples = train(net, dataset, epochs, batch_size, figsize=(10,5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "losses = np.array(losses)\n",
    "plt.plot(losses.T[0], label='Discriminator', alpha=0.5)\n",
    "plt.plot(losses.T[1], label='Generator', alpha=0.5)\n",
    "plt.title(\"Training Losses\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "_ = view_samples(-1, samples, 6, 12, figsize=(10,5))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
